"""
错别字纠错的Agent
"""

from nl2sql.model.llm import LLM

typo_prompt = """你是一个错别字纠错专家，专门负责把用户输入问题中的错别字改写成正确的内容。但是请仅限于修改错别字！
用户打错字可能是输入法的问题，请结合音译或别称来改写，如果没错，请不要动用户的问题。
用户问题: {query}
改写后的问题:"""


class TypoFixAgent:
    def __init__(self, llm: LLM):
        self.llm = llm
        self.typo_prompt = typo_prompt

    def fix_typo(self, query: str) -> str:
        prompt = self.typo_prompt.format(query=query)
        new_query: str = self.llm.chat(prompt=prompt)
        return new_query


if __name__ == '__main__':
    model_name = "qwen-max"
    api_key = "sk-34f5c792513c423f90b404d28b070f1f"
    base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"

    llm = LLM(model_name=model_name,
              api_key=api_key,
              base_url=base_url)

    tfa = TypoFixAgent(llm=llm)

    new = tfa.fix_typo(query='查一下各个gay都的客座率情况')
    print(new)
